package wei.xinya;

import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import org.apache.commons.lang3.RandomStringUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.http.HttpEntity;
import org.apache.http.HttpStatus;
import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.client.methods.HttpPost;
import org.apache.http.entity.ContentType;
import org.apache.http.entity.StringEntity;
import org.apache.http.entity.mime.MultipartEntityBuilder;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.HttpClients;
import org.apache.http.util.EntityUtils;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.Map;

public class CompreFace {

    public static void main(String[] args) throws Exception {
        if (args.length < 1) {
            System.out.println("参数异常，请输入图片目录！");
            return;
        }

        // 图片目录
        String faceDir = args[0];
        System.out.println("图片目录：" + faceDir);

        File faceDirFile = new File(faceDir);
        if (!faceDirFile.exists() || !faceDirFile.isDirectory()) {
            System.out.println("图片目录为空，程序退出！");
            System.exit(2);
        }

        LinkedList<File> files;

        while (true) {
            MyFileVisitor myFileVisitor = new MyFileVisitor(".jpg");
            Files.walkFileTree(faceDirFile.toPath(), myFileVisitor);

            files = myFileVisitor.getFiles();
            if (files.isEmpty()) {
                Thread.sleep(2000);
                continue;
            }

            for (File file : files) {
                if (!file.exists() || file.isDirectory()) {
                    continue;
                }

                String name = file.getName();
                if (!name.endsWith("BACKGROUND.jpg") && file.delete()) {
                    System.out.println("文件删除：" + file.getPath());
                    continue;
                }

                // 生成新文件名
                String randomStr = name.split("_")[2] + "_" + RandomStringUtils.randomAlphabetic(3);

                // 获取文件字节数组
                byte[] bytes = Files.readAllBytes(file.toPath());
                // 获取图片信息
                BufferedImage image = ImageIO.read(file);

                if (file.delete()) {
                    System.out.println("文件删除：" + file.getPath());
                }

                // 人像检测
                String res = detect(name, bytes);
                if (StringUtils.isBlank(res)) {
                    continue;
                }

                JSONObject jsonObject = JSONObject.parseObject(res);
                if (!jsonObject.containsKey("result")) {
                    continue;
                }

                JSONArray result = jsonObject.getJSONArray("result");
                if (result.isEmpty()) {
                    continue;
                }

                System.out.println("发现人像");

                for (int i = 0; i < result.size(); i++) {
                    JSONObject json = result.getJSONObject(i);
                    JSONObject age = json.getJSONObject("age");
                    JSONObject gender = json.getJSONObject("gender");
                    JSONObject box = json.getJSONObject("box");
                    if (age.getDoubleValue("probability") < 0.60
                            || gender.getDoubleValue("probability") < 0.70
                            || box.getDoubleValue("probability") < 0.80) {
                        System.out.println("人像质量低，不归档");
                        continue;
                    }

                    int x_max = box.getIntValue("x_max");
                    int y_max = box.getIntValue("y_max");
                    int x_min = box.getIntValue("x_min");
                    int y_min = box.getIntValue("y_min");

                    int width = x_max - x_min;
                    int height = y_max - y_min;
                    int init_x = x_min - (width / 2);
                    int init_y = y_min - (height / 2);

                    BufferedImage subImage;

                    try {
                        // 人像框截取
                        subImage = image.getSubimage(init_x, init_y, 2 * width, 2 * height);
                    } catch (Exception e) {
                        subImage = image.getSubimage(x_min, y_min, width, height);
                    }

                    // 人像小图文件
                    String fileName = randomStr + "_" + i + ".jpg";
                    ByteArrayOutputStream out = new ByteArrayOutputStream();
                    ImageIO.write(subImage, "jpg", out);

                    // 获取文件字节数组
                    bytes = out.toByteArray();

                    // 人脸身份识别
                    String recognize = recognize(fileName, bytes);
                    if (StringUtils.isBlank(recognize)) {
                        continue;
                    }

                    jsonObject = JSONObject.parseObject(recognize);
                    if (!jsonObject.containsKey("result")) {
                        continue;
                    }

                    result = jsonObject.getJSONArray("result");
                    if (result.isEmpty()) {
                        continue;
                    }

                    jsonObject = result.getJSONObject(0);
                    if (!jsonObject.containsKey("subjects")) {
                        continue;
                    }

                    JSONArray subjects = jsonObject.getJSONArray("subjects");
                    if (subjects.isEmpty()) {
                        // 添加人像库
                        jsonObject = new JSONObject();
                        jsonObject.put("subject", randomStr);
                        String subject = subjects(jsonObject.toString());
                        if (StringUtils.isBlank(subject)) {
                            System.out.println("人像库添加失败");
                            continue;
                        }

                        System.out.println("人像库：" + subject + " 添加成功");

                        // 添加人像到人像库
                        String faces = faces(randomStr, fileName, bytes);
                        if (StringUtils.isBlank(faces)) {
                            System.out.println("添加人像到人像库失败");
                            continue;
                        }

                        System.out.println("人像：" + faces + "添加成功");
                        continue;
                    }

                    jsonObject = subjects.getJSONObject(0);
                    if (jsonObject.getFloatValue("similarity") > 0.98) {
                        String subject = jsonObject.getString("subject");
                        System.out.println("发现：" + subject);

                        // 添加人像到人像库
                        String faces = faces(subject, fileName, bytes);
                        if (StringUtils.isBlank(faces)) {
                            System.out.println("添加人像到人像库失败");
                            continue;
                        }

                        System.out.println("人像：" + faces + "添加成功");
                        continue;
                    }

                    // 添加人像库
                    jsonObject = new JSONObject();
                    jsonObject.put("subject", randomStr);
                    String subject = subjects(jsonObject.toString());
                    if (StringUtils.isBlank(subject)) {
                        System.out.println("人像库添加失败");
                        continue;
                    }

                    System.out.println("人像库：" + subject + "添加成功");

                    // 添加人像到人像库
                    String faces = faces(randomStr, fileName, bytes);
                    if (StringUtils.isBlank(faces)) {
                        System.out.println("添加人像到人像库失败");
                        continue;
                    }

                    System.out.println("人像：" + faces + "添加成功");
                }
            }
        }
    }

    public static String faces(String subject, String fileName, byte[] bytes) {
        String url = "http://192.168.77.50:8000/api/v1/recognition/faces?subject=" + subject;

        HashMap<String, String> headers = new HashMap<>();
        headers.put("x-api-key", "afc8ca5a-6f72-4b65-9dd7-9a8472dafe7e");

        HttpEntity entity = MultipartEntityBuilder.create()
                .addBinaryBody("file", bytes, ContentType.IMAGE_JPEG, fileName).build();

        return doPost(url, headers, entity);
    }

    public static String subjects(String load) {
        String url = "http://192.168.77.50:8000/api/v1/recognition/subjects";

        HashMap<String, String> headers = new HashMap<>();
        headers.put("Content-Type", "application/json");
        headers.put("x-api-key", "afc8ca5a-6f72-4b65-9dd7-9a8472dafe7e");

        StringEntity entity = new StringEntity(load, StandardCharsets.UTF_8);

        return doPost(url, headers, entity);
    }

    public static String recognize(String fileName, byte[] bytes) {
        String url = "http://192.168.77.50:8000/api/v1/recognition/recognize?face_plugins=landmarks,gender,age";

        HashMap<String, String> headers = new HashMap<>();
        headers.put("x-api-key", "afc8ca5a-6f72-4b65-9dd7-9a8472dafe7e");

        HttpEntity entity = MultipartEntityBuilder.create()
                .addBinaryBody("file", bytes, ContentType.IMAGE_JPEG, fileName).build();

        return doPost(url, headers, entity);
    }

    public static String detect(String fileName, byte[] bytes) {
        String url = "http://192.168.77.50:8000/api/v1/detection/detect?face_plugins=landmarks,gender,age";

        HashMap<String, String> headers = new HashMap<>();
        headers.put("x-api-key", "54328958-cc36-4ce9-bc87-77e2e3cc9cf6");

        HttpEntity entity = MultipartEntityBuilder.create()
                .addBinaryBody("file", bytes, ContentType.IMAGE_JPEG, fileName).build();

        return doPost(url, headers, entity);
    }

    private static String doPost(String url, Map<String, String> headers, HttpEntity httpEntity) {
        String responseStr = null;
        CloseableHttpClient httpClient = null;
        HttpPost httpPost = new HttpPost(url);

        if (headers != null && !headers.isEmpty()) {
            for (Map.Entry<String, String> entry : headers.entrySet()) {
                httpPost.addHeader(entry.getKey(), entry.getValue());
            }
        }

        try {
            httpPost.setEntity(httpEntity);
        } catch (Exception e) {
            e.printStackTrace();
            return null;
        }

        try {
            httpClient = HttpClients.createDefault();
            try (CloseableHttpResponse response = httpClient.execute(httpPost)) {
                HttpEntity entity = response.getEntity();
                responseStr = EntityUtils.toString(entity, "UTF-8");

                int statusCode = response.getStatusLine().getStatusCode();
                if (statusCode != HttpStatus.SC_OK && statusCode != HttpStatus.SC_CREATED) {
                    return null;
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        } finally {
            if (httpClient != null) {
                try {
                    httpClient.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }

        return responseStr;
    }

}
